Step 1: Download a dataset

Download the Guns dataset from Roboflow's public databse.

Select YOLOv5 from the download formats and copy your unique terminal command into the cell below.

! curl -L "https://public.roboflow.com/ds/xXg040ThIk?key=mM2BVJ4kkZ" > data.zip \
  && unzip -q data.zip -d downloads \
  && rm data.zip
  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                 Dload  Upload   Total   Spent    Left  Speed
100   901  100   901    0     0   3575      0 --:--:-- --:--:-- --:--:--  3575
100 58.4M  100 58.4M    0     0  70.8M      0 --:--:-- --:--:-- --:--:-- 70.8M

Step 2: Explore the dataset

! pip install -q opencv-python-headless==4.5.4.60 fiftyone
import fiftyone as fo

data_path = "downloads/export/images"
labels_path = "downloads/export/labels"
classes = ["pistol"]

# Import dataset by explicitly providing paths to the source media and labels
dataset = fo.Dataset.from_dir(
    dataset_type=fo.types.YOLOv4Dataset,
    data_path=data_path,
    label_field="ground_truth",
    labels_path=labels_path,
    classes=classes,
)

# View data
session = fo.launch_app(dataset)

image.png

Step 3: Split the dataset into train/test/val subsets

import fiftyone.utils.splits as fous
import os

def split_dataset(dataset, split_dict, output_path):

  fous.random_split(dataset, split_dict, seed=42)

  for tag in list(split_dict.keys()):
    view = dataset.match_tags(tag)
    # Export each split separately
    dir = os.path.join(output_path, tag)

    # Export the dataset
    view.export(
        export_dir=dir,
        dataset_type=fo.types.YOLOv4Dataset,
        label_field="ground_truth",
    )


split_dataset(
    dataset=dataset, 
    split_dict={"train": 0.8, "test": 0.1, "val": 0.1}, 
    output_path="dataset/guns/"
)
 100% |███████████████| 2377/2377 [2.7s elapsed, 0s remaining, 859.1 samples/s]      
 100% |█████████████████| 297/297 [355.2ms elapsed, 0s remaining, 836.1 samples/s]      
 100% |█████████████████| 297/297 [397.3ms elapsed, 0s remaining, 747.5 samples/s]      
content = """# Train/val/test sets
path: /content/dataset/guns  # dataset root dir
train: train/data  # train images (relative to 'path')
test: test/data  # test images (relative to 'path')
val: val/data  # val images (relative to 'path')

# classes
nc: 1 # number of classes
names: ['pistol']"""

with open("pistols.yml", "w") as f:
  f.write(content)
! rm -r downloads

Step 4: Install DeepSparse and SparseML libraries and train our model

! pip install -q deepsparse[yolo] sparseml[torchvision]

Grab the transfer learning reciepe from github

! curl -LO "https://raw.githubusercontent.com/neuralmagic/sparseml/main/integrations/ultralytics-yolov5/recipes/yolov5.transfer_learn_pruned.md"
  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                 Dload  Upload   Total   Spent    Left  Speed
100  3084  100  3084    0     0  16146      0 --:--:-- --:--:-- --:--:-- 16146

Train the model using transfer learning on yolov5s (small) and the downloaded reciepe.

The batch-size = -1 parameter turns on autobatch.

! sparseml.yolov5.train \
  --data pistols.yml \
  --cfg models_v5.0/yolov5s.yaml \
  --weights zoo:cv/detection/yolov5-s/pytorch/ultralytics/coco/pruned-aggressive_96 \
  --hyp data/hyps/hyp.finetune.yaml \
  --batch-size -1 \
  --imgsz 416 \
  --recipe yolov5.transfer_learn_pruned.md
train: weights=zoo:cv/detection/yolov5-s/pytorch/ultralytics/coco/pruned-aggressive_96, cfg=models_v5.0/yolov5s.yaml, data=pistols.yml, hyp=data/hyps/hyp.finetune.yaml, epochs=300, batch_size=-1, imgsz=416, rect=False, resume=False, nosave=False, noval=False, noautoanchor=False, evolve=None, bucket=, cache=None, image_weights=False, device=, multi_scale=False, single_cls=False, optimizer=SGD, sync_bn=False, workers=8, project=yolov5_runs/train, name=exp, exist_ok=False, quad=False, cos_lr=False, label_smoothing=0.0, patience=0, freeze=[0], save_period=-1, local_rank=-1, entity=None, upload_dataset=False, bbox_interval=-1, artifact_alias=latest, recipe=yolov5.transfer_learn_pruned.md, disable_ema=False, max_train_steps=-1, max_eval_steps=-1, one_shot=False, num_export_samples=0
github: skipping check (not a git repository), for updates see https://github.com/ultralytics/yolov5
fatal: not a git repository (or any of the parent directories): .git
YOLOv5 🚀 2022-6-30 torch 1.9.1+cu102 CUDA:0 (Tesla T4, 15110MiB)

hyperparameters: lr0=0.0032, lrf=0.12, momentum=0.843, weight_decay=0.00036, warmup_epochs=2.0, warmup_momentum=0.5, warmup_bias_lr=0.05, box=0.0296, cls=0.243, cls_pw=0.631, obj=0.301, obj_pw=0.911, iou_t=0.2, anchor_t=2.91, fl_gamma=0.0, hsv_h=0.0138, hsv_s=0.664, hsv_v=0.464, degrees=0.373, translate=0.245, scale=0.898, shear=0.602, perspective=0.0, flipud=0.00856, fliplr=0.5, mosaic=1.0, mixup=0.243, copy_paste=0.0
Weights & Biases: run 'pip install wandb' to automatically track and visualize YOLOv5 🚀 runs (RECOMMENDED)
TensorBoard: Start with 'tensorboard --logdir yolov5_runs/train', view at http://localhost:6006/
Obtaining new sparse zoo credentials token
Getting signed url for c13e55cb-dd6c-4492-a079-8986af0b65e6/model.pt
Downloading model file model.pt to /root/.cache/sparsezoo/c13e55cb-dd6c-4492-a079-8986af0b65e6/pytorch/model.pt
downloading...: 100% 14.1M/14.1M [00:00<00:00, 37.0MB/s]
Overriding model.yaml nc=80 with nc=1

                 from  n    params  module                                  arguments                     
  0                -1  1      3520  yolov5.models.common.Focus              [3, 32, 3]                    
  1                -1  1     18560  yolov5.models.common.Conv               [32, 64, 3, 2]                
  2                -1  1     18816  yolov5.models.common.C3                 [64, 64, 1]                   
  3                -1  1     73984  yolov5.models.common.Conv               [64, 128, 3, 2]               
  4                -1  3    156928  yolov5.models.common.C3                 [128, 128, 3]                 
  5                -1  1    295424  yolov5.models.common.Conv               [128, 256, 3, 2]              
  6                -1  3    625152  yolov5.models.common.C3                 [256, 256, 3]                 
  7                -1  1   1180672  yolov5.models.common.Conv               [256, 512, 3, 2]              
  8                -1  1    656896  yolov5.models.common.SPP                [512, 512, [5, 9, 13]]        
  9                -1  1   1182720  yolov5.models.common.C3                 [512, 512, 1, False]          
 10                -1  1    131584  yolov5.models.common.Conv               [512, 256, 1, 1]              
 11                -1  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          
 12           [-1, 6]  1         0  yolov5.models.common.Concat             [1]                           
 13                -1  1    361984  yolov5.models.common.C3                 [512, 256, 1, False]          
 14                -1  1     33024  yolov5.models.common.Conv               [256, 128, 1, 1]              
 15                -1  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          
 16           [-1, 4]  1         0  yolov5.models.common.Concat             [1]                           
 17                -1  1     90880  yolov5.models.common.C3                 [256, 128, 1, False]          
 18                -1  1    147712  yolov5.models.common.Conv               [128, 128, 3, 2]              
 19          [-1, 14]  1         0  yolov5.models.common.Concat             [1]                           
 20                -1  1    296448  yolov5.models.common.C3                 [256, 256, 1, False]          
 21                -1  1    590336  yolov5.models.common.Conv               [256, 256, 3, 2]              
 22          [-1, 10]  1         0  yolov5.models.common.Concat             [1]                           
 23                -1  1   1182720  yolov5.models.common.C3                 [512, 512, 1, False]          
 24      [17, 20, 23]  1     16182  yolov5.models.yolo.Detect               [1, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]], [128, 256, 512]]
overriding activations in model to Hardswish
YOLOv5s summary: 283 layers, 7063542 parameters, 7063542 gradients, 16.5 GFLOPs

Transferred 354/361 items from /root/.cache/sparsezoo/c13e55cb-dd6c-4492-a079-8986af0b65e6/pytorch/model.pt
AutoBatch: Computing optimal batch size for --imgsz 416
AutoBatch: CUDA:0 (Tesla T4) 14.76G total, 0.06G reserved, 0.05G allocated, 14.64G free
      Params      GFLOPs  GPU_mem (GB)  forward (ms) backward (ms)                   input                  output
     7063542       6.961         0.201         16.78         14.15        (1, 3, 416, 416)                    list
     7063542       13.92         0.310         17.35         15.45        (2, 3, 416, 416)                    list
     7063542       27.85         0.581         18.58         18.79        (4, 3, 416, 416)                    list
     7063542       55.69         1.126         19.02         26.23        (8, 3, 416, 416)                    list
     7063542       111.4         2.223         32.87         50.45       (16, 3, 416, 416)                    list
AutoBatch: Using batch-size 96 for CUDA:0 13.28G/14.76G (90%)
Scaled weight_decay = 0.00054
optimizer: SGD with parameter groups 59 weight (no decay), 62 weight, 62 bias
albumentations: version 1.0.3 required by YOLOv5, but version 0.1.12 is currently installed
train: Scanning '/content/dataset/guns/train/data' images and labels...2377 found, 0 missing, 0 empty, 0 corrupt: 100% 2377/2377 [00:01<00:00, 1379.73it/s]
train: New cache created: /content/dataset/guns/train/data.cache
val: Scanning '/content/dataset/guns/val/data' images and labels...297 found, 0 missing, 0 empty, 0 corrupt: 100% 297/297 [00:00<00:00, 634.83it/s]
val: New cache created: /content/dataset/guns/val/data.cache
Plotting labels to yolov5_runs/train/exp2/labels.jpg... 

AutoAnchor: 2.95 anchors/target, 1.000 Best Possible Recall (BPR). Current anchors are a good fit to dataset ✅
Image sizes 416 train, 416 val
Using 2 dataloader workers
Logging results to yolov5_runs/train/exp2
Starting training for 300 epochs...
Disabling LR scheduler, managing LR using SparseML recipe
Overriding number of epochs from SparseML manager to 50

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      0/49     12.2G   0.07143  0.005311         0       243       416: 100% 25/25 [00:35<00:00,  1.43s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:04<00:00,  2.16s/it]
                 all        297        357     0.0179      0.185     0.0106    0.00189

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      1/49     14.3G    0.0591  0.005631         0       237       416: 100% 25/25 [00:32<00:00,  1.30s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00,  1.95s/it]
                 all        297        357       0.11      0.199     0.0727     0.0206

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      2/49     12.3G    0.0511  0.005889         0       186       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00,  1.61s/it]
                 all        297        357      0.356      0.401      0.275     0.0831

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      3/49     12.9G   0.04455  0.006083         0       195       416: 100% 25/25 [00:33<00:00,  1.33s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00,  1.53s/it]
                 all        297        357      0.419      0.473      0.406      0.132

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      4/49     12.9G   0.03806  0.006098         0       233       416: 100% 25/25 [00:32<00:00,  1.32s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00,  1.50s/it]
                 all        297        357      0.555      0.605      0.598      0.259

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      5/49     12.9G   0.03427  0.005962         0       258       416: 100% 25/25 [00:32<00:00,  1.30s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.46s/it]
                 all        297        357      0.729      0.591      0.686      0.347

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      6/49     12.9G   0.03197   0.00576         0       207       416: 100% 25/25 [00:33<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.47s/it]
                 all        297        357      0.718      0.669      0.718      0.392

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      7/49     12.9G   0.03004  0.005625         0       254       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00,  1.52s/it]
                 all        297        357      0.749      0.669      0.737      0.411

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      8/49       13G   0.02922  0.005537         0       235       416: 100% 25/25 [00:34<00:00,  1.38s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.44s/it]
                 all        297        357      0.701      0.691      0.752      0.436

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      9/49       13G   0.02814  0.005131         0       196       416: 100% 25/25 [00:33<00:00,  1.35s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.47s/it]
                 all        297        357      0.717      0.709      0.772      0.447

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     10/49       13G   0.02756  0.005152         0       247       416: 100% 25/25 [00:34<00:00,  1.37s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.48s/it]
                 all        297        357      0.671      0.787       0.78      0.461

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     11/49       13G   0.02719  0.005132         0       227       416: 100% 25/25 [00:32<00:00,  1.30s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00,  1.51s/it]
                 all        297        357      0.847      0.666        0.8      0.463

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     12/49       13G   0.02633  0.004897         0       265       416: 100% 25/25 [00:32<00:00,  1.32s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.48s/it]
                 all        297        357      0.722      0.773      0.811      0.495

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     13/49       13G   0.02559  0.005043         0       196       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.41s/it]
                 all        297        357      0.691      0.814      0.791      0.469

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     14/49       13G   0.02537  0.004815         0       232       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.44s/it]
                 all        297        357      0.791      0.731      0.811      0.499

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     15/49       13G   0.02512  0.004849         0       230       416: 100% 25/25 [00:34<00:00,  1.37s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.45s/it]
                 all        297        357      0.765      0.739       0.81      0.483

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     16/49       13G   0.02504  0.004759         0       206       416: 100% 25/25 [00:33<00:00,  1.33s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.47s/it]
                 all        297        357      0.737      0.779       0.82      0.512

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     17/49       13G   0.02469  0.004673         0       267       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.41s/it]
                 all        297        357      0.802       0.72      0.822      0.512

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     18/49       13G   0.02447  0.004662         0       231       416: 100% 25/25 [00:33<00:00,  1.35s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.46s/it]
                 all        297        357      0.778      0.759      0.824      0.522

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     19/49       13G   0.02424  0.004754         0       235       416: 100% 25/25 [00:33<00:00,  1.35s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.39s/it]
                 all        297        357      0.785      0.759      0.827      0.533

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     20/49       13G   0.02374  0.004643         0       213       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.42s/it]
                 all        297        357      0.857      0.695      0.823      0.517

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     21/49       13G   0.02402  0.004533         0       209       416: 100% 25/25 [00:34<00:00,  1.37s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.46s/it]
                 all        297        357      0.796      0.754      0.816      0.523

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     22/49     13.6G    0.0235  0.004528         0       236       416: 100% 25/25 [00:32<00:00,  1.31s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.40s/it]
                 all        297        357      0.777      0.782      0.833      0.545

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     23/49     13.6G   0.02366  0.004616         0       264       416: 100% 25/25 [00:33<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.44s/it]
                 all        297        357      0.816      0.748      0.832      0.531

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     24/49     13.6G   0.02322  0.004491         0       238       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.49s/it]
                 all        297        357      0.814      0.748      0.829      0.519

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     25/49     13.6G   0.02305  0.004363         0       213       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.40s/it]
                 all        297        357      0.768      0.773      0.835      0.536

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     26/49     13.6G   0.02326  0.004525         0       238       416: 100% 25/25 [00:33<00:00,  1.33s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.43s/it]
                 all        297        357      0.801      0.787      0.849      0.552

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     27/49     13.6G   0.02263  0.004288         0       250       416: 100% 25/25 [00:33<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.38s/it]
                 all        297        357      0.813       0.77      0.845      0.548

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     28/49     13.6G   0.02294  0.004502         0       245       416: 100% 25/25 [00:33<00:00,  1.33s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.46s/it]
                 all        297        357      0.821      0.762       0.84      0.546

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     29/49     13.6G   0.02249  0.004358         0       223       416: 100% 25/25 [00:33<00:00,  1.35s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.42s/it]
                 all        297        357      0.811       0.78      0.839      0.547

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     30/49     13.7G    0.0226  0.004365         0       234       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.40s/it]
                 all        297        357      0.832       0.75      0.837      0.545

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     31/49     13.7G   0.02216  0.004202         0       218       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.39s/it]
                 all        297        357      0.847      0.744      0.841      0.542

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     32/49     13.7G   0.02241  0.004421         0       258       416: 100% 25/25 [00:33<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.45s/it]
                 all        297        357      0.789      0.798      0.841      0.553

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     33/49     13.7G   0.02187  0.004324         0       183       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.42s/it]
                 all        297        357      0.798      0.798      0.847      0.551

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     34/49     13.7G   0.02235  0.004243         0       217       416: 100% 25/25 [00:32<00:00,  1.30s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.45s/it]
                 all        297        357      0.821      0.773      0.843       0.55

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     35/49     13.7G   0.02182  0.004277         0       234       416: 100% 25/25 [00:33<00:00,  1.35s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.42s/it]
                 all        297        357      0.804      0.796      0.849      0.559

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     36/49     13.7G   0.02203  0.004312         0       260       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.47s/it]
                 all        297        357      0.809      0.784      0.846      0.553

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     37/49     13.7G   0.02222   0.00428         0       200       416: 100% 25/25 [00:32<00:00,  1.31s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.43s/it]
                 all        297        357      0.806      0.804       0.85      0.557

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     38/49     13.7G   0.02202  0.004223         0       210       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.46s/it]
                 all        297        357      0.784      0.796      0.836      0.551

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     39/49     13.7G   0.02193  0.004199         0       256       416: 100% 25/25 [00:34<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.43s/it]
                 all        297        357      0.804      0.795      0.844      0.556

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     40/49     13.7G   0.02192  0.004291         0       246       416: 100% 25/25 [00:34<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.24s/it]
                 all        297        357      0.792        0.8      0.839      0.558

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     41/49     13.7G   0.02164  0.004269         0       235       416: 100% 25/25 [00:32<00:00,  1.31s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.47s/it]
                 all        297        357      0.783       0.81      0.842      0.561

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     42/49     13.7G   0.02207  0.004228         0       247       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.43s/it]
                 all        297        357       0.79      0.812      0.845      0.562

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     43/49     13.7G   0.02193  0.004293         0       263       416: 100% 25/25 [00:33<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00,  1.57s/it]
                 all        297        357      0.796      0.801       0.85      0.561

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     44/49     13.7G   0.02202  0.004254         0       242       416: 100% 25/25 [00:34<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.44s/it]
                 all        297        357       0.81      0.793      0.849      0.556

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     45/49     13.7G   0.02163  0.004189         0       209       416: 100% 25/25 [00:33<00:00,  1.33s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.48s/it]
                 all        297        357      0.791      0.807      0.849      0.557

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     46/49     13.7G   0.02168  0.004311         0       241       416: 100% 25/25 [00:34<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.44s/it]
                 all        297        357      0.801      0.801      0.844      0.556

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     47/49     13.7G   0.02202  0.004218         0       231       416: 100% 25/25 [00:33<00:00,  1.36s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.44s/it]
                 all        297        357      0.811      0.793      0.847      0.559

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     48/49     13.7G   0.02156  0.004261         0       230       416: 100% 25/25 [00:32<00:00,  1.32s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00,  1.51s/it]
                 all        297        357      0.835      0.765      0.842      0.556

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     49/49     13.7G   0.02169  0.004359         0       239       416: 100% 25/25 [00:33<00:00,  1.34s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00,  1.43s/it]
                 all        297        357      0.803       0.81      0.851      0.561

51 epochs completed in 0.526 hours.
Optimizer stripped from yolov5_runs/train/exp2/weights/last.pt, 14.4MB
Optimizer stripped from yolov5_runs/train/exp2/weights/best.pt, 14.4MB

Validating yolov5_runs/train/exp2/weights/best.pt...
Fusing layers... 
YOLOv5s summary: 224 layers, 7053910 parameters, 0 gradients, 16.3 GFLOPs
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00,  1.97s/it]
                 all        297        357      0.789      0.815      0.845      0.562
Results saved to yolov5_runs/train/exp2

Step 4: Export the Sparse Model to ONNX format for CPU inference

! sparseml.yolov5.export_onnx \
  --weights yolov5_runs/train/exp/weights/best.pt \
  --imgsz 416 \
  --dynamic
export: data=../usr/local/lib/python3.7/dist-packages/sparseml/yolov5/data/coco128.yaml, weights=['yolov5_runs/train/exp/weights/best.pt'], imgsz=[416], batch_size=1, device=cpu, half=False, inplace=False, train=False, optimize=False, int8=False, dynamic=True, simplify=False, opset=12, verbose=False, workspace=4, nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45, conf_thres=0.25, remove_grid=False, include=['torchscript', 'onnx']
export: data=../usr/local/lib/python3.7/dist-packages/sparseml/yolov5/data/coco128.yaml, weights=['yolov5_runs/train/exp/weights/best.pt'], imgsz=[416], batch_size=1, device=cpu, half=False, inplace=False, train=False, optimize=False, int8=False, dynamic=True, simplify=False, opset=12, verbose=False, workspace=4, nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45, conf_thres=0.25, remove_grid=False, include=['torchscript', 'onnx']
fatal: not a git repository (or any of the parent directories): .git
YOLOv5 🚀 2022-6-30 torch 1.9.1+cu102 CPU

Fusing layers... 
YOLOv5s summary: 224 layers, 7053910 parameters, 0 gradients, 16.3 GFLOPs

PyTorch: starting from yolov5_runs/train/exp/weights/best.pt with output shape (1, 10647, 6) (13.8 MB)

TorchScript: starting export with torch 1.9.1+cu102...
TorchScript: export success, saved as yolov5_runs/train/exp/weights/best.torchscript (27.2 MB)

ONNX: starting export with onnx 1.10.1...
ONNX: export success, saved as yolov5_runs/train/exp/weights/best.onnx (27.0 MB)

Export complete (7.26s)
Results saved to /content/yolov5_runs/train/exp/weights
Detect:          python detect.py --weights yolov5_runs/train/exp/weights/best.onnx
PyTorch Hub:     model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5_runs/train/exp/weights/best.onnx')
Validate:        python val.py --weights yolov5_runs/train/exp/weights/best.onnx
Visualize:       https://netron.app

Step 5: DeepSparse Engine model deployment

To validate the model on our test set in the DeepSparse Engine, run:

Note: this will be CPU only inference.

! sparseml.yolov5.val_onnx \
  --model_path yolov5_runs/train/exp/weights/best.onnx \
  --batch-size 33 \
  --imgsz 416 \
  --task test \
  --data pistols.yml
val_onnx: data=pistols.yml, model_path=yolov5_runs/train/exp/weights/best.onnx, batch_size=33, imgsz=416, conf_thres=0.001, iou_thres=0.6, task=test, device=, workers=8, augment=False, verbose=False, save_txt=False, save_hybrid=False, save_conf=False, save_json=False, project=../usr/local/lib/python3.7/dist-packages/sparseml/yolov5/runs/val, name=exp, exist_ok=False, half=False, dnn=False, engine=deepsparse, num_cores=None
val_onnx: data=pistols.yml, model_path=yolov5_runs/train/exp/weights/best.onnx, batch_size=33, imgsz=416, conf_thres=0.001, iou_thres=0.6, task=test, device=, workers=8, augment=False, verbose=False, save_txt=False, save_hybrid=False, save_conf=False, save_json=False, project=../usr/local/lib/python3.7/dist-packages/sparseml/yolov5/runs/val, name=exp, exist_ok=False, half=False, dnn=False, engine=deepsparse, num_cores=None
fatal: not a git repository (or any of the parent directories): .git
YOLOv5 🚀 2022-6-30 torch 1.9.1+cu102 CUDA:0 (Tesla T4, 15110MiB)

DeepSparse Engine, Copyright 2021-present / Neuralmagic, Inc. version: 1.0.0 (8eaddc24) (release) (optimized) (system=avx2, binary=avx2)
test: Scanning '/content/dataset/guns/test/data' images and labels...297 found, 0 missing, 0 empty, 0 corrupt: 100% 297/297 [00:00<00:00, 1260.64it/s]
test: New cache created: /content/dataset/guns/test/data.cache
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 9/9 [00:34<00:00,  3.88s/it]
/usr/local/lib/python3.7/dist-packages/yolov5/utils/metrics.py:74: FutureWarning: elementwise comparison failed; returning scalar instead, but in the future will perform elementwise comparison
  names = [v for k, v in names.items() if k in unique_classes]  # list: only classes that have data
                 all        297        351      0.918      0.764      0.886      0.578
Speed: 0.0ms pre-process, 116.0ms inference, 0.0ms NMS per image at shape (33, 3, 416, 416)
Results saved to ../usr/local/lib/python3.7/dist-packages/sparseml/yolov5/runs/val/exp2

Ste 6: Run the model on a youtube video using the Ultralytics YOLO library

We can also run the model trained with Neuralmagic CLI in using the Ultralytics YOLO library. Simply install the library and run inference using the *.pt model weights.

For this example we will use a scene from John Wick making Revolver to test the model.

! git clone https://github.com/ultralytics/yolov5
! pip install -q -r yolov5/requirements.txt
Cloning into 'yolov5'...
remote: Enumerating objects: 12351, done.
remote: Counting objects: 100% (16/16), done.
remote: Compressing objects: 100% (15/15), done.
remote: Total 12351 (delta 6), reused 4 (delta 1), pack-reused 12335
Receiving objects: 100% (12351/12351), 12.18 MiB | 31.10 MiB/s, done.
Resolving deltas: 100% (8520/8520), done.
! pip install -q -r yolov5/requirements.txt
!python yolov5/detect.py \
  --weights /content/yolov5_runs/train/exp/weights/best.pt \
  --source "https://www.youtube.com/watch?v=nB0A4pFIOA0" \
  --conf-thres=0.6 \
  --imgsz 416

Below are a few screenshots from the annotated video.

image.png

image.png

image.png